package com.example.mcpclient.controller;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.util.List;
import java.util.Map;

/**
 * 用到向量的
 *
 * @author xuguang
 */
@RestController
@RequestMapping("/vector")
@Slf4j
public class VectorController {

    @Autowired
    private VectorStore localVectorStore;

    @Autowired
    private VectorStore esVectorStore;

    /**
     * 添加文档到向量存储
     */
    @PostMapping("/add")
    public Mono<String> addDocument(@RequestBody DocumentRequest request) {
        return Mono.fromCallable(() -> {
            Document document = new Document(request.getContent(), request.getMetadata());
            localVectorStore.add(List.of(document));
            return "Document added successfully with id: " + document.getId();
        }).subscribeOn(Schedulers.boundedElastic());
    }

    /**
     * 根据ID删除文档
     */
    @PostMapping("/delete")
    public Mono<String> deleteDocument(@RequestBody DeleteRequest request) {
        return Mono.fromCallable(() -> {
            localVectorStore.delete(request.getIds());
            return "Documents deleted successfully";
        }).subscribeOn(Schedulers.boundedElastic());
    }


    /**
     * 相似性搜索
     */
    @PostMapping("/search")
    public Mono<List<Document>> searchDocuments(@RequestBody SearchRequestWrapper request) {
        return Mono.fromCallable(() -> {
            SearchRequest searchRequest = SearchRequest.builder().query(request.getQuery())
                    .topK(request.getTopK())
                    .similarityThreshold(request.getSimilarityThreshold()).build();
            return localVectorStore.similaritySearch(searchRequest);
        }).subscribeOn(Schedulers.boundedElastic());
    }


    /**
     * 添加文档到向量存储
     */
    @PostMapping("/esAdd")
    public Mono<String> esAdd(@RequestBody DocumentRequest request) {
        return Mono.fromCallable(() -> {
            Document document = new Document(request.getContent(), request.getMetadata());
            esVectorStore.add(List.of(document));
            return "Document added successfully with id: " + document.getId();
        }).subscribeOn(Schedulers.boundedElastic());
    }

    /**
     * 根据ID删除文档
     */
    @PostMapping("/esDelete")
    public Mono<String> esDelete(@RequestBody DeleteRequest request) {
        return Mono.fromCallable(() -> {
            esVectorStore.delete(request.getIds());
            return "Documents deleted successfully";
        }).subscribeOn(Schedulers.boundedElastic());
    }


    /**
     * 相似性搜索
     */
    @PostMapping("/esSearch")
    public Mono<List<Document>> esSearch(@RequestBody SearchRequestWrapper request) {
        return Mono.fromCallable(() -> {
            SearchRequest searchRequest = SearchRequest.builder().query(request.getQuery())
                    .topK(request.getTopK())
                    .similarityThreshold(request.getSimilarityThreshold()).build();
            return esVectorStore.similaritySearch(searchRequest);
        }).subscribeOn(Schedulers.boundedElastic());
    }

    @Data
    static class DocumentRequest {
        private String content;
        private Map<String, Object> metadata;
    }

    @Data
    static class DeleteRequest {
        private List<String> ids;
    }

    @Data
    static class SearchRequestWrapper {
        private String query;
        private int topK = 4;
        private double similarityThreshold = 0.3;
    }
}